#! /usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Author :   Liang Kang
@Contact:   gangkanli1219@163.com
@Time   :   2018/7/24 9:53
@Desc   :
"""
import hashlib
import os

import cv2
import numpy as np
import tensorflow as tf
from dltools.data import BaseRecordGenerator
from dltools.utils.io import read_voc_xml
from object_detection.utils import dataset_util


def _create_example(image, xml_path, shape, img_format='JPEG'):
    """
    创建一个 tensorflow example

    Parameters
    ----------
    image
    xml_path
    shape
    img_format

    Returns
    -------

    """
    boxes, labels, cls_name = [], [], []
    objects = read_voc_xml(xml_path, shape)
    for ob in objects:
        labels.append(ob['label'])
        boxes.append(ob['box'])
        cls_name.append(ob['name'].encode('utf8'))

    labels = np.asarray(labels) + 1
    labels = labels.tolist()
    boxes = np.asarray(boxes) / np.array([shape[:2] * 2])
    ymin, xmin, ymax, xmax = np.split(boxes, 4, 1)
    if isinstance(image, str):
        with tf.gfile.GFile(image, 'rb') as fid:
            encoded_jpg = fid.read()
    else:
        encoded_jpg = image.tobytes()
        img_format = 'RAW'
    key = hashlib.sha256(encoded_jpg).hexdigest()
    truncated, poses, difficult_obj = [], [], []
    example = tf.train.Example(
        features=tf.train.Features(feature={
            'image/height':
                dataset_util.int64_feature(shape[0]),
            'image/width':
                dataset_util.int64_feature(shape[1]),
            'image/filename':
                dataset_util.bytes_feature(
                    os.path.basename(xml_path).encode('utf8')),
            'image/source_id':
                dataset_util.bytes_feature(
                    os.path.basename(xml_path).encode('utf8')),
            'image/key/sha256':
                dataset_util.bytes_feature(key.encode('utf8')),
            'image/encoded':
                dataset_util.bytes_feature(encoded_jpg),
            'image/format':
                dataset_util.bytes_feature(img_format.encode('utf8')),
            'image/object/bbox/xmin':
                dataset_util.float_list_feature(xmin.tolist()),
            'image/object/bbox/xmax':
                dataset_util.float_list_feature(xmax.tolist()),
            'image/object/bbox/ymin':
                dataset_util.float_list_feature(ymin.tolist()),
            'image/object/bbox/ymax':
                dataset_util.float_list_feature(ymax.tolist()),
            'image/object/class/text':
                dataset_util.bytes_list_feature(cls_name),
            'image/object/class/label':
                dataset_util.int64_list_feature(labels),
            'image/object/difficult':
                dataset_util.int64_list_feature(difficult_obj),
            'image/object/truncated':
                dataset_util.int64_list_feature(truncated),
            'image/object/view':
                dataset_util.bytes_list_feature(poses),
        }))
    return example.SerializeToString()


class RecordGenerator(BaseRecordGenerator):
    """
    利用数据文件生成 tensorflow record 文件
    """

    def __init__(self, data, output, logger=None, display=10):
        """

        Parameters
        ----------
        data: 数据
        output: 输出文件
        logger: 日志对象
        display: 每转换多少文件显示一次
        """
        if logger is not None:
            logger = logger.getChild('RecordGenerator')
        super(RecordGenerator, self).__init__(data, output, display, logger)

    def _encode_data(self):
        image_path = self._buf_data['raw']['image']
        xml_path = self._buf_data['raw']['xml']
        image = cv2.imread(image_path)
        shape = image.shape
        self._buf_data['data'] = _create_example(image_path, xml_path, shape)

    def _write_data(self, writer):
        example = self._buf_data['data']
        writer.write(example)
